# This code is hosted on http://code.google.com/p/lenthorp/
# Freely available for use in applications, but should NOT be modified
# Email all comments to lenthorpresearch@gmail.com

from FiniteDifferencesEngineMod import *
from copy import deepcopy
from numpy import interp
from Tools import multiDimensionalLinearInterpolation
from math import pow

class ExplicitFiniteDifferencesEngine(FiniteDifferencesEngine):

    def initialise(self, relations, initials):
        self.nodeCoeffsToCalculate = []
        self.relations = relations
        for relation in self.relations:
            for varName in relation.variableNames:
                relation.varIndices.append(self.grid.findIndexInNamesOrder(varName))
        

        # calculate the nodes required ...
        #   first, find the 'end' node
        self.firstCoord = [0 for idx in range(len(initials))]
        self.distances = [0 for idx in range(len(initials))]
        for idx in range(len(initials)):
            gridIdx = self.grid.findIndexInNamesOrder(initials.keys()[idx])
            for jdx in range(self.grid.dimSizes[gridIdx]):
                if self.grid.g[gridIdx][jdx] >= initials.values()[idx]:
                    self.firstCoord[gridIdx] = jdx
                    self.distances[gridIdx] = self.grid.g[gridIdx][jdx] - initials.values()[idx]
                    break

        #   now move 'forwards', looping over all coords and doing special nodes last
        self.transientIndex
        coords = [0 for idx in range(self.grid.ndims)]
        variableIndices = [i for i in range(self.grid.ndims)]
        variableIndices.remove(self.transientIndex)

        isForwards = self.firstCoord[self.transientIndex] == self.grid.dimSizes[self.transientIndex] - 1
        for trans in range(self.firstCoord[self.transientIndex], -1 if isForwards else self.grid.dimSizes[self.transientIndex], -1 if isForwards else 1):
            coords[self.transientIndex] = trans
            saveCoords = []
            currentDimToChange = 0
            while currentDimToChange < len(variableIndices):
                if coords[variableIndices[currentDimToChange]] < self.grid.dimSizes[variableIndices[currentDimToChange]]:
                    if coords in self.BoundaryConditionNodes:
                        pass
                    elif coords in self.SpecialNodes:
                        saveCoords.append(deepcopy(coords))
                    else:
                        self.nodeCoeffsToCalculate.append(deepcopy(coords))
                    coords[variableIndices[currentDimToChange]] += 1
                else:
                    coords[variableIndices[currentDimToChange]] = 0
                    currentDimToChange += 1
                    if currentDimToChange == len(variableIndices):
                        break
                    coords[variableIndices[currentDimToChange]] += 1
                    if coords[variableIndices[currentDimToChange]] < self.grid.dimSizes[variableIndices[currentDimToChange]]:
                        currentDimToChange = 0
            self.nodeCoeffsToCalculate += saveCoords

        #   finally, reverse the whole thing
        self.nodeCoeffsToCalculate.reverse()

    def doCalculation(self):
        for coeff in self.nodeCoeffsToCalculate:
            val = 0.0
            
            if coeff in self.SpecialNodes:
                relations = self.SpecialRelations[self.SpecialNodes.index(coeff)]
            else:
                relations = self.relations
            for relation in relations:
                multiplier = 1.0
                if len(relation.variableNames) > 0:
                    u = deepcopy(coeff)
                    for idx in range(len(relation.variableNames)):
                        u[relation.varIndices[idx]] += relation.numberSteps[idx]
                    multiplier = self.grid.getData(u)
                val += relation.coeffFunc(self.grid.getDistancesToNextNode(coeff), self.grid.getVariableValues(coeff)) * multiplier
            if val < 0: # for debugging
                val = val
            self.grid.setData(coeff, val)

        # do the grid interpolation for each dimension 
        totalDist = self.grid.getDistancesToNextNode(self.firstCoord)
        nDims = len(self.firstCoord)
        vals = []
        dists = []
        leftCoord = deepcopy(self.firstCoord)
        for idx in range(nDims):
            leftCoord[idx] = 0 if leftCoord[idx]==0 else leftCoord[idx]-1
            dists.append([0.0, totalDist[idx]] if leftCoord[idx]==0 else [totalDist[idx] - self.distances[idx], self.distances[idx]])
        for idx in range(int(pow(2,nDims))):
            num = bin(idx).replace('0b','')
            while len(num) < nDims:
                num = '0' + num
            currentCoord = [0 for jdx in range(nDims)]
            for jdx in range(nDims):
                currentCoord[jdx] = leftCoord[jdx] + int(num[jdx])
            vals.append(self.grid.getData(currentCoord))
        valFinal = multiDimensionalLinearInterpolation(vals, dists)
        return valFinal
            